package com.rrd.tools.lang.staticextender.annotation;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.builder.EqualsBuilder;
import org.apache.commons.lang3.builder.HashCodeBuilder;
import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.LocalVariablesSorter;
import org.objectweb.asm.util.ASMifier;
import org.objectweb.asm.util.TraceClassVisitor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.rrd.tools.attach.LocalInstrumentationFactory;

/**
 * An rrd-attach-util based utility to add the ability to override static methods to java.
 * 
 * @see OverrideStatic
 * @see LocalInstrumentationFactory
 * @author erachitskiy
 *
 */
public class OverrideStatics {
	private static final Logger log = LoggerFactory.getLogger(OverrideStatics.class);
	/**
	 * Class that can be used to represent an internal method invokation and translate it
	 * to an actual {@link Method} instance.
	 * @author erachitskiy
	 *
	 */
	private static class MethodDescriptor{
		private final String methodName;
		private final String methodDesc;
		private final String methodOwner;
		public MethodDescriptor(String methodName, String methodDesc, String methodOwner) {
			super();
			this.methodName = methodName;
			this.methodDesc = methodDesc;
			this.methodOwner = methodOwner;
		}		
		@Override
		public String toString() {
			return methodOwner+"."+methodName+methodDesc;
		}
		@Override
		public boolean equals(Object obj) {
			return EqualsBuilder.reflectionEquals(this, obj);
		}
		@Override
		public int hashCode() {
			return HashCodeBuilder.reflectionHashCode(this);
		}		
	}
	/**
	 * ClassVisitor that can be used to detect an overriden method provided a descriptor of an overriding method
	 * @author erachitskiy
	 *
	 */
	private static class OverridenMethodLocator extends ClassVisitor{
		String superClassName;
		String myName;
		MethodDescriptor matchingMethod;
		MethodDescriptor overridingMethod;
		public OverridenMethodLocator(MethodDescriptor overridingMethod) {
			super(Opcodes.ASM4);
			this.overridingMethod=overridingMethod;
		}
		@Override
		public void visit(int version, int access, String name, String signature, String superName,
				String[] interfaces) {
			myName = name;
			superClassName = superName;
			super.visit(version, access, name, signature, superName, interfaces);
		}
		@Override
		public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
			if(name.equals(overridingMethod.methodName)&&desc.equals(overridingMethod.methodDesc)){
				matchingMethod = new MethodDescriptor(name, desc, myName);
			}
			return super.visitMethod(access, name, desc, signature, exceptions);
		}
		public MethodDescriptor getMatchingMethod() {
			return matchingMethod;
		}
		public String getSuperClassName() {
			return superClassName;
		}
	}
	/**
	 * ClassVisitor that looks for calls to any overriden methods and replaces them with calls to the
	 * {@link OverrideStatics#intercept(Object, String, String, String, Object...)} method.
	 * @author erachitskiy
	 *
	 */
	private static class MethodChanger extends ClassVisitor{
		String myClassName;
		public MethodChanger(ClassVisitor cv) {
			super(Opcodes.ASM4, cv);
		}
		
		@Override
		public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
			log.debug("visiting {}",name);
			myClassName = name;
			super.visit(version, access, name, signature, superName, interfaces);
		}
		@Override
		public MethodVisitor visitMethod(int access, final String methodName, final String methodDesc, String signature,
				String[] exceptions) {					
			/* first calculate maxStack and maxLocals via local sorter */
			final LocalVariablesSorter sorter = new LocalVariablesSorter(access, methodDesc, super.visitMethod(access,
					methodName, methodDesc, signature, exceptions));
			return new MethodVisitor(Opcodes.ASM4, sorter) {
				int maxStackDelta = 0;
				@Override
				public void visitMethodInsn(int opcode, String owner, String name, String desc) {						
					if (opcode == Opcodes.INVOKESTATIC && overridenMethods.containsKey(new MethodDescriptor(name, desc, owner))) {
						log.debug("Found a static invokation {} {} {} in {} {} {}",owner,name,desc,myClassName,methodName,methodDesc);
						/*
						 * at this point stack contains all arguments the method is
						 * going to be called with copy these into an array
						 */

						List<Integer> stackArgumentVariables = new LinkedList<Integer>();

						/*
						 * compute arguments length and create array of size of
						 * arguments
						 */
						Type[] methodArguments = Type.getArgumentTypes(desc);
						for (int i = 0; i < methodArguments.length; i++) {
							/*
							 * use the sorter to locate a new local variable index
							 * (and auto-increment maxLocals)
							 */
							int newLocalVariableNum = sorter.newLocal(methodArguments[i]);
							stackArgumentVariables.add(newLocalVariableNum);
							/* store each method argument into a local variable */
							mv.visitIntInsn(Opcodes.ASTORE, newLocalVariableNum);
						}

						/* load this */
						mv.visitVarInsn(Opcodes.ALOAD, 0);
						/* load owner */
						mv.visitLdcInsn(owner);
						/* load target static method name */
						mv.visitLdcInsn(name);
						/* load target method desc */
						mv.visitLdcInsn(desc);
						/* create an array */
						/* create a new arguments array */
						mv.visitIntInsn(Opcodes.SIPUSH, methodArguments.length);
						mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/Object");
						/* now pop items off the stack into the array */
						for (int i = 0; i < methodArguments.length; i++) {
							/* push arrayref onto stack */
							mv.visitInsn(Opcodes.DUP);
							/* push arary index */
							mv.visitIntInsn(Opcodes.SIPUSH, i);
							/* load the argument from local var */
							mv.visitIntInsn(Opcodes.ALOAD, stackArgumentVariables.get(i));
							/* store argument in array */
							mv.visitInsn(Opcodes.AASTORE);
						}

						/* call static method */
						mv.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getType(OverrideStatics.class).getInternalName() ,
								"intercept",
								"(Ljava/lang/Object;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/Object;");
						/* figure out if we need to perform a cast */
						Type returnType = Type.getReturnType(desc);	
						
						/* cast primitives if necessary */
						switch(returnType.getSort()){
							case Type.BOOLEAN:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Boolean");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Boolean", "booleanValue", "()Z");
								break;
							case Type.BYTE:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Byte");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Byte", "byteValue", "()B");
								break;
							case Type.CHAR:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Character");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Character", "charValue", "()C");
								break;
							case Type.SHORT:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Short");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Short", "shortValue", "()S");
								break;
							case Type.INT:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Integer");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Integer", "intValue", "()I");
								break;
							case Type.FLOAT:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Float");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Float", "floatValue", "()F");
								break;
							case Type.LONG:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Long");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Long", "longValue", "()J");
								break;
							case Type.DOUBLE:
								mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Double");
								mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Double", "doubleValue", "()D");
								break;
							case Type.ARRAY:
								mv.visitTypeInsn(Opcodes.CHECKCAST, returnType.getDescriptor());
								break;
							case Type.OBJECT:
								mv.visitTypeInsn(Opcodes.CHECKCAST, returnType.getInternalName());
								break;
							case Type.VOID:
								/* do nothing for voids! */
								break;
						}
																						
						/* increment max stack delta */
						maxStackDelta += 6+methodArguments.length;						
					} else {

						super.visitMethodInsn(opcode, owner, name, desc);
					}
				}

				@Override
				public void visitEnd() {
					if(maxStackDelta>0){
						log.debug("Rewrote {} {} {}",myClassName,methodName,methodDesc);
					}
					super.visitEnd();
				}
				@Override
				public void visitMaxs(int maxStack, int maxLocals) {
					super.visitMaxs(maxStack + maxStackDelta+50, maxLocals+50);
				}
			};

		}
		
	}
	private static class OverridesStaticClassFileTransformer implements ClassFileTransformer{
		public byte[] transform(final ClassLoader loader, final String internalClassName, Class<?> classBeingRedefined,
				ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {						
			/* find super classes */
			/* this has to happen here since when an annotation is detected, the containing classe's subclasses can not be resolved */
			Set<String> superClasses = getAllSuperClassInternalNames(new ClassReader(classfileBuffer));
			/* add self to set of super classes */
			superClasses.add(internalClassName);
			/* check if this class or any of its super classes have overriden methods */
			if(!Collections.disjoint(superClasses, internalClassNamesToTransform)){
				/* ignore already transformed classes */
				if(alreadyTransformedClasses.contains(internalClassName)){
					return classfileBuffer;
				}
				/* in case this is a subclass of a class with an override annotation, 
				 * need to add all of its inner classes to the transformation list */
				internalClassNamesToTransform.addAll(findAllInnerClasses(internalClassName));
				ClassReader cr = new ClassReader(classfileBuffer);
				log.debug("Attempting to re transform {} ({}b)",internalClassName,classfileBuffer.length);
				
				if(log.isDebugEnabled()){
					StringWriter sw = new StringWriter();
					TraceClassVisitor tcv = new TraceClassVisitor(null, new ASMifier(), new PrintWriter(sw));
					cr.accept(tcv, 0);					
					log.debug("Original class:\n{}\n",sw.toString());					
				}
				
				/* transform this class */				
				ClassWriter cw = new ClassWriter(0);
				MethodChanger mc = new MethodChanger(cw);
				cr.accept(mc, 0);				
				alreadyTransformedClasses.add(internalClassName);
				byte[] newClassfileBuffer = cw.toByteArray();
				log.debug("Obtained new bytecode for {} ({}b)",internalClassName,newClassfileBuffer.length);
				
				if(log.isDebugEnabled()){
					StringWriter sw = new StringWriter();
					TraceClassVisitor tcv = new TraceClassVisitor(null, new ASMifier(), new PrintWriter(sw));
					new ClassReader(newClassfileBuffer).accept(tcv, 0);					
					log.debug("New class:\n{}\n",sw.toString());					
				}
				
				return newClassfileBuffer;						
			}else{
				/* find all annotations and map them to methods */
				ClassReader cr = new ClassReader(classfileBuffer);
				cr.accept(new ClassVisitor(Opcodes.ASM4) {
					String superClassInternalName;
					@Override
					public void visit(int version, int access, String name, String signature, String superName,
							String[] interfaces) {
						superClassInternalName = superName;
						super.visit(version, access, name, signature, superName, interfaces);
					}
					@Override
					public MethodVisitor visitMethod(int access, final String methodName, final String methodDesc, String signature,
							String[] exceptions) {						
						return new MethodVisitor(Opcodes.ASM4) {
							@Override
							public AnnotationVisitor visitAnnotation(String annotationDesc, boolean visible) {							
								if (OverrideStatic.class.getName().equals(Type.getType(annotationDesc).getClassName())) {
									MethodDescriptor overridingMethod = new MethodDescriptor(methodName, methodDesc, internalClassName);
									/* iterate all parent classes and try to find overriden method */
									String curClassInternalName = superClassInternalName;
									Set<String> classesToTransform = new HashSet<String>();
									MethodDescriptor overridenMethod = null;
									while(curClassInternalName!=null&&!"java/lang/Object".equals(curClassInternalName)){
										try{
											classesToTransform.add(curClassInternalName);
											classesToTransform.addAll(findAllInnerClasses(curClassInternalName));
											ClassReader cr = new ClassReader(curClassInternalName);
											OverridenMethodLocator l = new OverridenMethodLocator(overridingMethod);
											cr.accept(l, 0);
											if(overridenMethod == null && l.getMatchingMethod()!=null){
												overridenMethod = l.getMatchingMethod();
											}
											curClassInternalName = l.getSuperClassName();
										}catch(Exception e){
											log.error("Could not locate overriden method for {}:{}",overridingMethod,e.getMessage());
										}
									}
									/* if an overriding method was found - go ahead and add to the list of internal classes to transform */
									if(overridenMethod!=null){
										internalClassNamesToTransform.addAll(classesToTransform);
										overridenMethods.put(overridenMethod, overridingMethod);
										classesWithOverrideMethods.add(internalClassName);
									}
									
								}
								return super.visitAnnotation(annotationDesc, visible);
							}
						};
					}
				}, 0);
				return classfileBuffer;
			}			
		}		
	}
	/* a map of overriden methods to their overriding methods */
	private static final Map<MethodDescriptor,MethodDescriptor> overridenMethods = new HashMap<MethodDescriptor, MethodDescriptor>();
	private static final Set<String> internalClassNamesToTransform = new HashSet<String>();
	private static final Set<String> alreadyTransformedClasses = new HashSet<String>();
	/* a set of classes with override methods that actually have a valid method they override */
	private static final Set<String> classesWithOverrideMethods = new HashSet<String>();
	private static OverridesStaticClassFileTransformer transformer = null;
	private static boolean enabled = false;
	
	private static Set<String> findAllInnerClasses(String className){
		Set<String> innerClasses = new HashSet<String>();
		findAllInnerClassesRecursive(className, innerClasses);
		return innerClasses;
	}
	private static void findAllInnerClassesRecursive(final String className,final Set<String> innerClasses){		
		try {
			ClassReader cr = new ClassReader(className);
			cr.accept(new ClassVisitor(Opcodes.ASM4) {
				public void visitInnerClass(String name, String outerName, String innerName, int access) {
					if(!innerClasses.contains(name)){
						innerClasses.add(name);
						findAllInnerClassesRecursive(name, innerClasses);
					}
				};
			}, 0);
		} catch (IOException e) {					
			throw new RuntimeException(e.getMessage(),e);			
		}		
	}

	
	private static void retransformClasses(Class<?> ... classes){
		try {
			LocalInstrumentationFactory.getLocalInstrumentation().retransformClasses(classes);
		} catch (Throwable e) {
			/* ignore */
			/* TODO: should potentially log errors here */
		}	
	}
	/**
	 * Enable static overrides.  All methods annotated with {@link OverrideStatic} will override calls to super class
	 * methods with the same name for all instances of these methods' enclosing classes. 
	 *  For example:
	 * {@code
	 * 	class Base{
	 * 		String callStaticMethod(){
	 * 			return aStaticMethod();
	 * 		}
	 * 		static String aStaticMethod(){
	 * 			return "Hello from Base";
	 * 		}
	 *  }
	 *  
	 *  class Subclass extends Base{
	 *  	@OverridesStatic
	 * 		static String aStaticMethod(){
	 * 			return "Hello from Subclass";
	 * 		}
	 *  }
	 * }
	 * in the above example, once {@link OverrideStatics#enable()} is called,
	 *  {@code new Subclass().callStaticMethod()} will return "Hello from Subclass"
	 *   whereas {@code new Base().callStaticMethod()} will return "Hello from Base"
	 */
	public static synchronized void enable(){
		enabled = true;
		if(transformer==null){					
			transformer = new OverridesStaticClassFileTransformer();
			LocalInstrumentationFactory.getLocalInstrumentation().addTransformer(transformer,true);
		}
		
		for(Class<?> loadedClass:LocalInstrumentationFactory.getLocalInstrumentation().getAllLoadedClasses()){			
			retransformClasses(loadedClass);						
		}
		log.debug("Added {} and transformed {}, with map {} after first pass",internalClassNamesToTransform,alreadyTransformedClasses,overridenMethods);
		Set<String> classesThatNeedTransformation = new HashSet<String>(internalClassNamesToTransform);
		classesThatNeedTransformation.removeAll(alreadyTransformedClasses);
		
		for(String needsTransofmration:classesThatNeedTransformation){
			try {				
				/* TODO: should resolve class loaders for these types if possible instead of calling Class.forName */
				retransformClasses(Class.forName(Type.getObjectType(needsTransofmration).getClassName()));
			} catch (ClassNotFoundException e) {
				log.warn("Could not transform {}:{}",needsTransofmration,e.getMessage());
			}
		}	
		
	}
	/**
	 * Disable static overrides.  All methods annotated with {@link OverrideStatic} will no longer override calls to super 
	 * class methods with the same name for all instances of these methods' enclosing classes.
	 */
	public static synchronized void disable(){		
		if(transformer!=null){
			LocalInstrumentationFactory.getLocalInstrumentation().removeTransformer(transformer);
		}
		transformer = null;		
		/* clear already transformed classes in case enable is called again after disabling */
		alreadyTransformedClasses.clear();
		enabled = false;
		/* TODO: enabling/disabling could potentially be a lot quicker if 
		 * 	* the transformer was registered permanently
		 *  * alreadyTransformedClasses wasn't reset (but somehow monitored)
		 *  * on enabling, only classes that haven't been transformed are retransformed
		 */
	}
	public static Object intercept(Object instance, String target, String name, String desc, Object... args) {

		log.debug("Intercepted static method call from {} to {} {} {} with args {}",instance,target,name,desc,args);		
		/* see if method should be redirected */			
		MethodDescriptor interceptedMethod = new MethodDescriptor(name, desc, target);
		Method toInvoke = null;						
		ClassLoader instanceLoader = instance.getClass().getClassLoader();
		if(enabled /* override statics is enabled */
				&&isInstanceAnExtendingClass(instance) /* instance should have its invocations redirected */
				&&overridenMethods.containsKey(interceptedMethod) /* this is indeed an overriden method call */				
				){
			/* should redirect */			
			toInvoke = getMethodForDescriptor(overridenMethods.get(interceptedMethod), instanceLoader);
			log.debug("Invoking override method ({})",toInvoke);
		}else{
			/* should call original method */
			log.debug("Invoking original method");
			toInvoke = getMethodForDescriptor(interceptedMethod,instanceLoader);
		}		
		try {
			return toInvoke.invoke(null, args);
		} catch (Exception e) {
			throw new RuntimeException(e.getMessage(),e);
		}
	}
	private static boolean isInstanceAnExtendingClass(Object instance){		
		
		/* check if instance is an instance of an actual class with an override method */
		if(classesWithOverrideMethods.contains(Type.getType(instance.getClass()).getInternalName())){
			return true;
		}
		/* check if instance has an enclosing type that is a class with an override method */
		for(Class<?> enclosingType:getEnclosingTypesForInstance(instance)){
			Class<?> instanceClass = enclosingType;
			do{
				log.debug("Checking if enclosing type {} is overriding class",instanceClass);
				if(classesWithOverrideMethods.contains(Type.getType(instanceClass).getInternalName())){
					return true;
				}
				instanceClass=instanceClass.getSuperclass();
			}while(instanceClass!=null);						
		}
		/* check if a superclass of instance is of a type with an override method */
		Class<?> instanceClass = instance.getClass();
		while((instanceClass=instanceClass.getSuperclass())!=null){
			if(classesWithOverrideMethods.contains(Type.getType(instanceClass).getInternalName())){
				return true;
			}
		}
		
		
		return false;
	}
	private static Set<Class<?>> getEnclosingTypesForInstance(Object instance){
		Field[] fields = instance.getClass().getDeclaredFields();
		Set<Class<?>> ret = new HashSet<Class<?>>();
		for(Field field:fields){
			/* TODO: classes can technically legally declare fields named this$XYZ */
			if(StringUtils.startsWith(field.getName(), "this$")){
				field.setAccessible(true);
				try {					
					Object enclosingInstance = field.get(instance);
					ret.add(enclosingInstance.getClass());
					/* TODO: perhaps add a seen-before protection from stack overflow here */
					ret.addAll(getEnclosingTypesForInstance(enclosingInstance));
				} catch (Exception e) {
					throw new RuntimeException();
				}								
			}			
		}
		return ret;
	}
	
	private static Method getMethodForDescriptor(MethodDescriptor methodDescriptor,ClassLoader loader){			
		try{				
			Class<?> owner = loader.loadClass(Type.getObjectType(methodDescriptor.methodOwner).getClassName()); 						
			Type[] types = Type.getMethodType(methodDescriptor.methodDesc).getArgumentTypes();
			Class<?> args[] = new Class<?>[types.length];
			for(int i=0;i<types.length;i++){
				args[i]=Class.forName(types[i].getClassName());
			}
			Method method = owner.getDeclaredMethod(methodDescriptor.methodName, args);
			method.setAccessible(true);
			return method;
		}catch(Exception e){								
			throw new RuntimeException(e.getMessage(),e);
		}
	}
	
	static Set<String> getAllSuperClassInternalNames(ClassReader cr){		
		final Set<String> internalSuperNames = new HashSet<String>();
		cr.accept(new ClassVisitor(Opcodes.ASM4) {			
			@Override
			public void visit(int version, int access, String name, String signature, String superName,
					String[] interfaces) {
				if(superName!=null){
					internalSuperNames.add(superName);
					/* read super class */				
					try {
						internalSuperNames.addAll(getAllSuperClassInternalNames(new ClassReader(superName)));
					} catch (IOException e) {
						log.error("Could not read super class {}:{}",superName,e.getMessage());
					}
				}
				super.visit(version, access, name, signature, superName, interfaces);
			}
		}, 0);		
		return internalSuperNames;
	}
}
